# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import collections
import functools
import logging
import math
import typing

import networkx

from coin.experimental.xguo.fruit.util.fee import get_fee_rate
from coin.exchange.base.kr_rest.currency import NativeCurrency
from coin.exchange.base.order_gateway import OrderSide
from coin.experimental.xguo.fruit.base.single_product import SingleProductPriceCalculator


class Order(object):
  def __init__(
      self,
      *,
      product=None,
      price=None,
      qty=None,
      source_qty=None,
      target_qty=None,
      order_side=None,
      edge_info=None,
  ):
    self.product = product
    self.price = price
    self.qty = qty
    self.source_qty = source_qty
    self.target_qty = target_qty
    self.order_side = order_side
    self.edge_info = edge_info


class PriceGraph(object):
  def __init__(self, products, pivots=None, ignore_fee=False):
    self._products = products
    self._pivots = pivots
    self._graph = networkx.DiGraph()
    self._price_calculators = {}
    self._timestamp = None
    self._all_path = []

    # initialize, graph topology does not change, only edge info will be
    # updated by book
    self._init_graph(ignore_fee)

  def _init_graph(self, ignore_fee):
    for product in self._products:
      base = product.base
      quote = product.quote
      price_calculator = SingleProductPriceCalculator(product)
      if ignore_fee:
        fee_bp = None
      else:
        fee_bp = get_fee_rate(product)

      buy_edge_info = EdgeInfo(
          price_calculator=price_calculator,
          order_side=OrderSide.BUY,
          fee_bp=fee_bp,
      )
      sell_edge_info = EdgeInfo(
          price_calculator=price_calculator,
          order_side=OrderSide.SELL,
          fee_bp=fee_bp,
      )
      self._graph.add_edge(quote, base, edge_info=buy_edge_info)
      self._graph.add_edge(base, quote, edge_info=sell_edge_info)
      self._price_calculators[product] = price_calculator
      price_calculator.set_price_det_notional(0.0001)

    for path in networkx.simple_cycles(self._graph):
      if len(path) > 2:
        path_info = PathInfo(path=path, graph=self)
        print(path_info)
        if self._pivots is not None:
          path_info = path_info.rotate(self._pivots)
        self._all_path.append(path_info)

  def is_ready(self):
    return all(calculator.is_ready() for calculator in self._price_calculators.values())

  def update(self, product, book):
    self._timestamp = book.timestamp
    base = product.base
    quote = product.quote
    price_calculator = self._price_calculators[product]
    price_calculator.update(book)
    if price_calculator.is_ready():
      self.get_edge_info(quote, base).update()
      self.get_edge_info(base, quote).update()
      for path in self.get_relevant_path(product):
        path.calculate_path_weight()

  def get_edge_info(self, source, target):
    info = self._graph.edges[source, target]['edge_info']
    return info

  def get_weight(self, source, target):
    edge_info = self.get_edge_info(source, target)
    return edge_info.weight

  @functools.lru_cache()
  def get_relevant_path(self, product):
    res = []
    for path in self._all_path:
      if product in path.products:
        res.append(path)
    return res

  @property
  def paths(self):
    return self._all_path

  def arb_chance(self, profit_bp):
    for path in self._all_path:
      path.calculate_path_weight()
    sorted_all_path = sorted(self._all_path, key=lambda x: x.weight, reverse=True)
    if sorted_all_path[0].weight > profit_bp * 1e-4:
      return sorted_all_path[0]
    else:
      return None

  def gen_order(self):
    all_orders = collections.defaultdict(dict)
    for path in self._all_path:
      orders = path.gen_mm_orders(edge_bp=edge_bp)
      for order in orders:
        all_orders[order.product][order.order_side] = order
    return all_orders


class EdgeInfo(object):
  def __init__(
      self,
      *,
      price_calculator: SingleProductPriceCalculator,
      order_side: OrderSide,
      fee_bp=None,
  ):
    self.price_calculator = price_calculator
    self.order_side = order_side
    self.price = None

    if fee_bp is None:
      fee_bp = [0.0, 0.0]
    elif isinstance(fee_bp, (int, float)):
      fee_bp = [fee_bp, fee_bp]

    self.maker_fee = fee_bp[0] * 1e-4
    self.taker_fee = fee_bp[1] * 1e-4
    self.fee = self.taker_fee

    self.weight = -99999999
    self._logger = logging.getLogger(__name__)

  @property
  def product(self):
    return self.price_calculator.product

  @property
  def timestamp(self):
    return self.price_calculator.timestamp

  def is_ready(self):
    return self.price_calculator.is_ready()

  def _update_buy_edge(self):
    price = self.price_calculator.ask0
    self.price = price
    self.weight = math.log(1. / price) - self.fee

  def _update_sell_edge(self):
    price = self.price_calculator.bid0
    self.price = price
    self.weight = math.log(price) - self.fee

  def update(self):
    if self.order_side == OrderSide.BUY:
      self._update_buy_edge()
    elif self.order_side == OrderSide.SELL:
      self._update_sell_edge()
    else:
      raise ValueError('Invalid Order Side: %s', self.order_side)

  def weight_to_price(self, weight):
    if self.order_side == OrderSide.BUY:
      return math.exp(-weight)
    elif self.order_side == OrderSide.SELL:
      return math.exp(weight)
    else:
      raise ValueError('Invalid order side!')

  def gen_order(self, source_qty):
    # TODO: add fee mode? currency assume always deduct from target
    target_qty = math.exp(math.log(source_qty) + self.weight)
    if self.order_side == OrderSide.BUY:
      qty = target_qty
    elif self.order_side == OrderSide.SELL:
      qty = source_qty
    else:
      raise ValueError('Unknown order side!')

    order = Order(
        product=self.product,
        price=self.price,
        qty=qty,
        source_qty=source_qty,
        target_qty=target_qty,
        order_side=self.order_side,
    )
    return order


class PathInfo(object):
  def __init__(
      self,
      *,
      path: typing.List[NativeCurrency],
      graph: PriceGraph,
  ):
    self._graph = graph
    self._path = path
    self._weight = -99999999

  def __len__(self):
    return len(self._path)

  def __str__(self):
    node_str_list = []
    for node in self._path:
      node_str_list.append(node.full_symbol)
    str_form = ' -> '.join(node_str_list) + ' weight: %.02f' % (self._weight * 1e4)

    weight_list = []
    for edge_info in self.edge_infos:
      weight_list.append(edge_info.weight)

    prices = [edge_info.price for edge_info in self.edge_infos]
    str_form += (' ' + str(weight_list) + ' ' + str(prices))
    return str_form

  @property
  def weight(self):
    return self._weight

  @property
  @functools.lru_cache(maxsize=1)
  def products(self):
    products = []
    path_len = len(self)
    for n in range(path_len):
      source = self._path[n]
      target = self._path[(n + 1) % path_len]
      edge_info = self._graph.get_edge_info(source, target)
      products.append(edge_info.product)
    return products

  @property
  @functools.lru_cache(maxsize=1)
  def currencies(self):
    currencies = set()
    for product in self.products:
      currencies.add(product.base)
      currencies.add(product.quote)
    return list(currencies)

  @property
  def orders(self):
    orders = []
    path_len = len(self)
    for n in range(path_len):
      source = self._path[n]
      target = self._path[(n + 1) % path_len]
      edge_info = self._graph.get_edge_info(source, target)
      order = Order(
          product=edge_info.product,
          price=edge_info.price,
          qty=None,
          order_side=edge_info.order_side,
          edge_info=edge_info,
      )
      orders.append(order)
    return orders

  @property
  def edge_infos(self):
    edge_infos = []
    path_len = len(self)
    for n in range(path_len):
      source = self._path[n]
      target = self._path[(n + 1) % path_len]
      edge_info = self._graph.get_edge_info(source, target)
      edge_infos.append(edge_info)
    return edge_infos

  def rotate(self, pivots):
    path_len = len(self._path)
    path = self._path + self._path
    for pivot in pivots:
      idx = path.index(pivot)
      new_path = path[idx:idx + path_len]
      return type(self)(path=new_path, graph=self._graph)
    raise ValueError('Invalid pivots! %s' % str(self))

  def calculate_path_weight(self):
    path_len = len(self)
    weight = 0
    for n in range(len(self)):
      source = self._path[n]
      target = self._path[(n + 1) % path_len]
      weight += self._graph.get_weight(source, target)
    self._weight = weight
    return weight

  def gen_mm_orders(self, edge_bp):
    orders = [None] * len(self.edge_infos)
    for idx in range(len(self.edge_infos)):
      edge = self.edge_infos[idx]
      weight = 0
      for n, edge_info in enumerate(self.edge_infos):
        if n != idx:
          weight -= self.edge_infos[n].weight
      price = edge.weight_to_price(weight + edge_bp / 1e4)
      orders[idx] = Order(
          product=edge.product,
          price=price,
          order_side=edge.order_side,
          qty=None,
          edge_info=edge,
      )
    return orders


def main():
  PriceGraph()


if __name__ == '__main__':
  main()
